Skip to content

Add opt-in idle model unloading to free VRAM - #889

Open
michael-borck wants to merge 2 commits into
jamiepine:mainfrom
michael-borck:feat/idle-model-unload
Open

Add opt-in idle model unloading to free VRAM#889
michael-borck wants to merge 2 commits into
jamiepine:mainfrom
michael-borck:feat/idle-model-unload

Conversation

@michael-borck

@michael-borck michael-borck commented Jul 13, 2026

Copy link
Copy Markdown

Closes #595 (the VRAM half of it — this adds the automatic trigger, not the manual load/unload UI button).

Problem

On a headless/server deployment there's no way for models to unload automatically. They load on demand at generate time, but the only things that ever unload them are an explicit POST /models/unload or process shutdown (app.py). There's no idle path.

So an instance that served one request an hour ago is still pinning its weights — ~5.6GB for qwen-tts-1.7B, which is most of an 8GB card. Anything else sharing that GPU is starved until you restart the container or remember to call the endpoint by hand.

Change

An asyncio reaper that unloads models after a period of inactivity.

  • Off by default. Gated on VOICEBOX_IDLE_UNLOAD_SECONDS; 0/unset disables it entirely, so existing behaviour is unchanged for anyone who doesn't opt in.
  • Reuses the existing unload primitives — the same tts / transcribe / llm unload calls already used on shutdown, with per-model failure isolation so one failure can't block the others.
  • Can't race in-flight work. "Idle" is read from the serial generation queue's own state (_queued_generation_ids / _running_generation_tasks), not inferred externally, so a model is never pulled out from under a queued or running generation. A long generation continually re-arms the countdown.
  • The next request reloads the model transparently — it just pays the cold-start it already pays on first use.

Three files: a new backend/services/idle_unload.py, two lines in app.py to start it, and a mark_activity() call in the generation worker's finally block (the one place every generation exits through, success/failure/cancel alike).

Testing

Manual, on an RTX 2060 Super (8GB), Docker/CUDA, qwen-tts-1.7B:

phase GPU memory
baseline, nothing loaded 442 MiB
model loaded, idle 5166 MiB
after reaper fires 874 MiB

Reclaims ~4.3GB of resident model weights. Verified separately that:

  • the reaper does not fire during a running generation (watched it hold off across a ~200s job that peaked at 6504 MiB, with a 120s idle timeout that would otherwise have been eligible mid-job);
  • the model reloads correctly on the next request after an auto-unload, and that generation completes (status=completed, no errors);
  • with the env var unset, the reaper never starts.

I did not run Black across the touched files — app.py has pre-existing formatting drift (the all_origins CORS block) and reformatting it would add unrelated churn to this diff. The new file and my task_queue.py edit are both Black-clean on their own.


Disclosure: this patch was written with AI assistance (Claude), as was this description. It's tested end-to-end on my own build and the numbers above are real measurements from my hardware, not generated. If AI-assisted contributions aren't welcome here I understand — happy for this to be closed, and it stays a local patch on my build.

Summary by CodeRabbit

  • New Features
    • Added optional automatic unloading of speech and language models after a configurable idle period.
    • Unloading is coordinated with generation activity to avoid interrupting queued or in-progress work.
    • Idle activity tracking now resets after each generation attempt, including failures and cancellations.

Models load on demand at generate time but are only unloaded explicitly
(POST /models/unload) or at shutdown, so an idle instance pins its weights
indefinitely -- ~5.6GB for qwen-tts-1.7B, most of an 8GB card.

Adds an asyncio reaper that unloads after a period of inactivity, gated on
VOICEBOX_IDLE_UNLOAD_SECONDS and off by default. Idleness is read from the
serial generation queue's own state so a model is never pulled out from
under queued or in-flight work; the next request reloads it transparently.

Refs jamiepine#595
@coderabbitai

coderabbitai Bot commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 7688c2d1-0122-4f1c-8aed-7bb62cd15c25

📥 Commits

Reviewing files that changed from the base of the PR and between 47e63c3 and 13005f4.

📒 Files selected for processing (1)
  • backend/services/idle_unload.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • backend/services/idle_unload.py

📝 Walkthrough

Walkthrough

The backend adds an optional idle reaper that tracks generation activity and unloads TTS, Whisper, and LLM models after a configured idle period. It starts during application startup and resets its timer after each generation cycle.

Changes

Idle model unloading

Layer / File(s) Summary
Idle reaper and model unloading
backend/services/idle_unload.py
Configures idle timing, tracks queue activity, detects inactivity, and unloads each model type in a background thread with per-model error isolation.
Application startup integration
backend/app.py, backend/services/idle_unload.py
Starts the idle-unload service after queue initialization when the configured timeout is positive.
Generation activity reset
backend/services/task_queue.py
Resets the idle countdown after every generation worker iteration, including success, failure, and cancellation paths.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant AppStartup
  participant GenerationQueue
  participant IdleReaper
  participant ModelBackends

  AppStartup->>GenerationQueue: init_queue()
  AppStartup->>IdleReaper: start()
  GenerationQueue->>IdleReaper: mark_activity() after each work cycle
  IdleReaper->>GenerationQueue: check queued and running generations
  IdleReaper->>ModelBackends: unload TTS, Whisper, and LLM after timeout
Loading
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Linked Issues check ⚠️ Warning #595 asks for manual load/unload controls and loaded-state visibility, but this PR only adds automatic idle unloading. Add explicit Load/Unload controls and loaded/not-loaded status, or retitle the PR to reflect the automatic idle-unload scope.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title is concise and accurately describes the main change: opt-in idle unloading to free VRAM.
Out of Scope Changes check ✅ Passed The changes stay focused on idle unloading, startup wiring, and activity tracking, with no obvious unrelated additions.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@backend/services/idle_unload.py`:
- Around line 76-82: Update the backend-loaded guard in _reaper so it considers
all backends cleared by _unload_all, including TTS, Whisper/STT, and LLM, before
skipping an unload cycle. Either remove the guard or reuse the corresponding
backend loaded-state checks, while preserving the existing exception-safe reaper
behavior.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 6f4b4425-a983-4b08-8fd9-998861fa4425

📥 Commits

Reviewing files that changed from the base of the PR and between f2cf2a7 and 47e63c3.

📒 Files selected for processing (3)
  • backend/app.py
  • backend/services/idle_unload.py
  • backend/services/task_queue.py

Comment thread backend/services/idle_unload.py Outdated
michael-borck added a commit to michael-borck/loco-puente that referenced this pull request Jul 13, 2026
Upstream voicebox never unloads models once loaded, so an idle instance
pins ~4.3GB of GPU 1 indefinitely (most of the 2060 Super's 8GB).

Adds VoiceboxConfig.build_ref (owner/repo@ref) so the service can build
from a fork carrying the idle-unload patch, and defaults the timeout to
600s. Pointed at michael-borck/voicebox@feat/idle-model-unload pending
jamiepine/voicebox#889; drop build_ref from puente.yml once that merges.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The reaper's guard only consulted the TTS backend, but _unload_all()
frees Whisper and the LLM too. A session that dictates or chats without
ever generating speech leaves TTS unloaded, so the guard short-circuited
every cycle and those models stayed resident indefinitely — the exact
leak this feature is meant to fix.

Check all three backends, isolating each getter so one failing backend
can't mask a loaded sibling. On error assume loaded rather than skipping:
_unload_all() already isolates its own per-model failures, so attempting
the unload is safer than pinning weights forever, and the exception is
now logged instead of swallowed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add manual model load/unload control instead of auto-loading models on generate

1 participant